/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.fnexecution.state;



import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.StateRequest;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnApi.StateResponse;

import com.bff.gaia.unified.model.fnexecution.v1.UnifiedFnStateGrpc;

import com.bff.gaia.unified.runners.fnexecution.FnService;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.stub.ServerCallStreamObserver;

import com.bff.gaia.unified.vendor.grpc.v1p13p1.io.grpc.stub.StreamObserver;



import java.util.concurrent.*;



import static com.bff.gaia.unified.vendor.guava.com.google.common.base.Throwables.getStackTraceAsString;



/** An implementation of the Unified Fn State service. */

public class GrpcStateService extends UnifiedFnStateGrpc.UnifiedFnStateImplBase

    implements StateDelegator, FnService {

  /** Create a new {@link GrpcStateService}. */

  public static GrpcStateService create() {

    return new GrpcStateService();

  }



  private final ConcurrentLinkedQueue<Inbound> clients;

  private final ConcurrentMap<String, StateRequestHandler> requestHandlers;



  private GrpcStateService() {

    this.requestHandlers = new ConcurrentHashMap<>();

    this.clients = new ConcurrentLinkedQueue<>();

  }



  @Override

  public void close() throws Exception {

    Exception thrown = null;

    for (Inbound inbound : clients) {

      try {

        // the call may be cancelled because the sdk harness hung up

        // (we terminate the environment before terminating the service endpoints)

        if (inbound.outboundObserver instanceof ServerCallStreamObserver) {

          if (((ServerCallStreamObserver) inbound.outboundObserver).isCancelled()) {

            // skip to avoid call already closed exception

            continue;

          }

        }

        inbound.outboundObserver.onCompleted();

      } catch (Exception t) {

        if (thrown == null) {

          thrown = t;

        } else {

          thrown.addSuppressed(t);

        }

      }

    }

    if (thrown != null) {

      throw thrown;

    }

  }



  @Override

  public StreamObserver<StateRequest> state(StreamObserver<StateResponse> responseObserver) {

    Inbound rval = new Inbound(responseObserver);

    clients.add(rval);

    return rval;

  }



  @Override

  public StateDelegator.Registration registerForProcessBundleInstructionId(

      String processBundleInstructionId, StateRequestHandler handler) {

    requestHandlers.putIfAbsent(processBundleInstructionId, handler);

    return new Registration(processBundleInstructionId);

  }



  private class Registration implements StateDelegator.Registration {

    private final String processBundleInstructionId;



    private Registration(String processBundleInstructionId) {

      this.processBundleInstructionId = processBundleInstructionId;

    }



    @Override

    public void deregister() {

      requestHandlers.remove(processBundleInstructionId);

    }



    @Override

    public void abort() {

      deregister();

      // TODO: Abort in-flight state requests. Flag this processBundleInstructionId as a fail.

    }

  }



  /**

   * An inbound {@link StreamObserver} which delegates requests to registered handlers.

   *

   * <p>Is only threadsafe if the outbound observer is threadsafe.

   *

   * <p>TODO: Handle when the client indicates completion or an error on the inbound stream and

   * there are pending requests.

   */

  private class Inbound implements StreamObserver<StateRequest> {

    private final StreamObserver<StateResponse> outboundObserver;



    Inbound(StreamObserver<StateResponse> outboundObserver) {

      this.outboundObserver = outboundObserver;

    }



    @Override

    public void onNext(StateRequest request) {

      StateRequestHandler handler =

          requestHandlers.getOrDefault(request.getInstructionReference(), this::handlerNotFound);

      try {

        CompletionStage<StateResponse.Builder> result = handler.handle(request);

        result.whenComplete(

            (StateResponse.Builder responseBuilder, Throwable t) ->

                // note that this is threadsafe if and only if outboundObserver is threadsafe.

                outboundObserver.onNext(

                    t == null

                        ? responseBuilder.setId(request.getId()).build()

                        : createErrorResponse(request.getId(), t)));

      } catch (Exception e) {

        outboundObserver.onNext(createErrorResponse(request.getId(), e));

      }

    }



    @Override

    public void onError(Throwable t) {

      outboundObserver.onCompleted();

    }



    @Override

    public void onCompleted() {

      outboundObserver.onCompleted();

    }



    private CompletionStage<StateResponse.Builder> handlerNotFound(StateRequest request) {

      CompletableFuture<StateResponse.Builder> result = new CompletableFuture<>();

      result.complete(

          StateResponse.newBuilder()

              .setError(

                  String.format(

                      "Unknown process bundle instruction id '%s'",

                      request.getInstructionReference())));

      return result;

    }



    private StateResponse createErrorResponse(String id, Throwable t) {

      return StateResponse.newBuilder().setId(id).setError(getStackTraceAsString(t)).build();

    }

  }

}